import os
import argparse
import models
import datasets

file_dir = os.path.dirname(__file__)  # the directory that options.py resides in
model_names = sorted(name for name in models.__dict__
                     if name.islower() and not name.startswith("__"))
dataset_names = sorted(name for name in datasets.__all__)

class ManyFlow:
    def __init__(self):
        self.parser = argparse.ArgumentParser(description="ManyFlow options")

        #PATHS
        self.parser.add_argument("--data_path",
                                 type=str,
                                 help="path to the training data",
                                 default=os.path.join(file_dir, "Data"))
        self.parser.add_argument("--log_dir",
                                 type=str,
                                 help="log directory",
                                 default=os.path.join(file_dir, "tmp"))

        # TRAINING options
        self.parser.add_argument('--dataset', metavar='DATASET', default='flying_chairs',
                            choices=dataset_names,
                            help='dataset type : ' +
                                 ' | '.join(dataset_names))
        self.parser.add_argument('-s', '--split-file', default="datasets/split", type=str,
                           help='test-val split file')
        self.parser.add_argument('--split-value', default=0.8, type=float,
                           help='test-val split proportion between 0 (only test) and 1 (only train), '
                                'will be overwritten if a split file is set')
        self.parser.add_argument('--sparse', action='store_true',
                            help='look for NaNs in target flow when computing EPE, avoid if flow is garantied to be dense,'
                                 'automatically seleted when choosing a KITTIdataset')
        self.parser.add_argument('--pretrained', dest='pretrained',  default="model/FlowNet.pth",
                            help='path to pre-trained model')
        self.parser.add_argument('--no-date', action='store_true',
                            help='don\'t append date timestamp to folder')


        #Abliation study
        self.parser.add_argument('--arch', '-a', metavar='ARCH', default='flownets',
                                 choices=model_names,
                                 help='model architecture, overwritten if pretrained is specified: ' +
                                      ' | '.join(model_names))
        self.parser.add_argument('--solver', default='adam', choices=['adam', 'sgd'],
                                 help='solver algorithms')
        self.parser.add_argument('--multiscale-weights', '-w', default=[0.005, 0.01, 0.02, 0.08, 0.32], type=float,
                                 nargs=5,
                                 help='training weight for each scale, from highest resolution (flow2) to lowest (flow6)',
                                 metavar=('W2', 'W3', 'W4', 'W5', 'W6'))
        self.parser.add_argument('--div-flow', default=20,
                                 help='value by which flow will be divided. Original value is 20 but 1 with batchNorm gives good results')


        # OPTIMIZATION options
        self.parser.add_argument('--freeze_teacher',
                                 action="store_true",
                                 help="If set, freeze the weights of the single frame teacher"
                                      " network and pose network.")
        self.parser.add_argument("--freeze_teacher_epoch",
                                 type=int,
                                 default=250,
                                 help="Sets the epoch number at which to freeze the teacher"
                                      "network and the pose network.")
        self.parser.add_argument('-b', '--batch-size', default=8, type=int,
                                 metavar='N', help='mini-batch size')
        self.parser.add_argument('--num_epochs', default=300, type=int, metavar='N',
                            help='number of total epochs to run')
        self.parser.add_argument('--lr', '--learning-rate', default=1e-4, type=float,
                                 metavar='LR', help='initial learning rate')
        self.parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                                 help='momentum for sgd, alpha parameter for adam')
        self.parser.add_argument('--beta', default=0.999, type=float, metavar='M',
                                 help='beta parameter for adam')
        self.parser.add_argument('--weight-decay', '--wd', default=4e-4, type=float,
                                 metavar='W', help='weight decay')
        self.parser.add_argument('--bias-decay', default=0, type=float,
                                 metavar='B', help='bias decay')
        self.parser.add_argument('--milestones', default=[100, 200], metavar='N', nargs='*',
                            help='epochs at which learning rate is divided by 2')

        # LOGGING options
        self.parser.add_argument("--load_teacher_weights_folder",
                                 default="model",
                                 type=str)
        self.parser.add_argument("--load_weights_folder",
                                 type=str,
                                 help="name of model to load")
        self.parser.add_argument("--models_to_load",
                                 nargs="+",
                                 type=str,
                                 help="models to load",
                                 default=["Flow_costvolume"])
        self.parser.add_argument("--log_frequency",
                                 type=int,
                                 help="number of batches between each tensorboard log",
                                 default=250)
        self.parser.add_argument("--save_frequency",
                                 type=int,
                                 help="number of epochs between each save",
                                 default=30)
        self.parser.add_argument('--print-freq', '-p', default=10, type=int,
                                 metavar='N', help='print frequency')


        # SYSTEM options
        self.parser.add_argument("--no_cuda",
                                 help="if set disables CUDA",
                                 action="store_true")
        self.parser.add_argument('-j', '--workers', default=12, type=int, metavar='N',
                                 help='number of data loading workers')

    def parse(self):
        self.options = self.parser.parse_args()
        return self.options